from PyCmpltrtok.common import sep
from neo4j import GraphDatabase
from python_nlp.kg.neo4j.config import NEO4J_CONFIG_AGR

confirm_location_suffix = ['省', '市', '区', '县', '镇', '乡', '村', '自治区', '自治州']

def confirm_location(driver, xname):
    xcandidates = []

    suffixes = ['', *confirm_location_suffix]

    with driver.session() as ss:
        # 依次加后缀，找到存在的节点并返回
        for suf in suffixes:
            cypher = "MATCH (n:Node{name:%r}) RETURN n.name" % (xname + suf, )
            record = ss.run(cypher)
            if record:
                for row in record:
                    xcandidates.append(row[0])

    return xcandidates



if '__main__' == __name__:

    def _main():
        # 驱动
        driver = GraphDatabase.driver(**NEO4J_CONFIG_AGR)

        while True:
            sep()
            print('请输入查询的地区：')
            xname = input().strip()
            xcandidates = confirm_location(driver, xname)
            print(xcandidates)

    _main()
